package edu.cornell.cs4740.qa.ner;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import edu.cornell.cs4740.qa.WordContainer;
import edu.stanford.nlp.ie.AbstractSequenceClassifier;
import edu.stanford.nlp.ie.crf.CRFClassifier;
import edu.stanford.nlp.ling.CoreAnnotations.AnswerAnnotation;
import edu.stanford.nlp.ling.CoreLabel;

public class NamedEntityRecognizer {

	// NAMED ENTITY CATEGORIES
	public static final String NULL_ENTITY = "0";
	public static final String PERSON_ENTITY = "PERSON";
	public static final String TIME_ENTITY = "TIME";
	public static final String ORG_ENTITY = "ORGANIZATION";
	public static final String LOCATION_ENTITY = "LOCATION";
	public static final String MONEY_ENTITY = "MONEY";
	public static final String PERCENT_ENTITY = "PERCENT";
	public static final String DATE_ENTITY = "DATE";

	AbstractSequenceClassifier<CoreLabel> classifier;

	@SuppressWarnings("unchecked")
	public NamedEntityRecognizer() {
		String serializedModel = "models/english.muc.7class.distsim.crf.ser.gz";

		classifier = CRFClassifier.getClassifierNoExceptions(serializedModel);
	}

	public String findNamedEntities(String s) {

		return classifier.classifyToString(s);
	}

	public String findNamedEntitiesXml(String s) {

		return classifier.classifyWithInlineXML(s);
	}

	public Map<String, String> getNamedEntityMap(String sentenceText) {
		Map<String, String> namedEntityMap = new HashMap<String, String>();

		List<List<CoreLabel>> cl = classifier.classify(sentenceText);

		String curPhrase = "";
		String lastAnnotation = null;

		for (List<CoreLabel> list : cl) {
			for (CoreLabel label : list) {
				String ansAnnotation = label.get(AnswerAnnotation.class);

				if (curPhrase.equals("")) {
					lastAnnotation = ansAnnotation;
					curPhrase = label.word();
				} else {
					if (ansAnnotation.equals(lastAnnotation)) {
						curPhrase += label.before();
						curPhrase += label.word();
					} else {
						// flush into the map
						namedEntityMap.put(curPhrase, lastAnnotation);


						lastAnnotation = ansAnnotation;
						curPhrase = label.word();

					}
				}
			}
		}

		if (!("".equals(curPhrase))) {
			namedEntityMap.put(curPhrase, lastAnnotation);
		}

		// for(int x = 0; x < namedEntityArray.length; x++) {
		// System.out.println(namedEntityArray[x]);
		// }

		return namedEntityMap;
	}

	@Deprecated
	public Map<String, String> getEntitiesInSentence(String sent){
		return this.getNamedEntityMap(sent);
	}
	
	// DOES NOT WORK AT ALL DO NOT USE
	@Deprecated
	public List<WordContainer> enrichWordContainersWithNamedEntities(List<WordContainer> wordContainers) {

		List<WordContainer> enrichedAndMergedWordContainers = new ArrayList<WordContainer>();

		String fullSentence = "";
		for (WordContainer wc : wordContainers) {
			fullSentence += wc.getWord();
			fullSentence += " ";
		}

		String namedEntitySentence = findNamedEntities(fullSentence);

		String[] namedEntityArray = namedEntitySentence.split("/0");

		for (int x = 0; x < namedEntityArray.length; x++) {
			System.out.println(namedEntityArray[x]);
		}

		return enrichedAndMergedWordContainers;

	}
}
